import os
import cv2
import json
import numpy as np
import time

class NumpyEncoder(json.JSONEncoder):
    """自定义 JSON 编码器，用于处理 NumPy 数据类型"""
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NumpyEncoder, self).default(obj)

class DatasetGenerator:
    def __init__(self, recognition, output_dir=".", slices_dir="slices", original_image_path=None):
        self.recognition = recognition
        self.output_dir = output_dir
        self.slices_dir = os.path.join(output_dir, slices_dir)
        self.original_image_path = original_image_path or recognition.imgPath
    
    def generate_dataset(self):
        """生成包含所有符号标注的数据集"""
        # 创建数据集字典
        dataset = {
            "image_path": self.original_image_path,
            "annotations": []
        }
        
        # 尝试获取符号列表
        symbols = []
        if hasattr(self.recognition, 'get_symbols') and callable(self.recognition.get_symbols):
            symbols = self.recognition.get_symbols()
            print(f"使用 get_symbols 方法，找到 {len(symbols)} 个符号")
        elif hasattr(self.recognition, 'symbols'):
            symbols = self.recognition.symbols
            print(f"使用 symbols 属性，找到 {len(symbols)} 个符号")
        elif hasattr(self.recognition, 'symbol_list'):
            symbols = self.recognition.symbol_list
            print(f"使用 symbol_list 属性，找到 {len(symbols)} 个符号")
        else:
            print("错误: 无法获取识别符号列表")
            return self.output_dir
        
        # 确保切片目录存在
        os.makedirs(self.slices_dir, exist_ok=True)
        
        # 获取图像尺寸
        if hasattr(self.recognition, 'thImg') and self.recognition.thImg is not None:
            img_height, img_width = self.recognition.thImg.shape[:2]
        else:
            img_height, img_width = 0, 0
        
        # 遍历所有识别结果
        for i, symbol in enumerate(symbols):
            # 跳过没有位置信息的符号
            if 'position' not in symbol or len(symbol['position']) != 4:
                print(f"警告: 符号 {i} 没有有效的位置信息")
                continue
                
            # 提取位置信息
            position = symbol['position']
            try:
                # 确保位置坐标是整数
                y_min = int(position[0])
                y_max = int(position[1])
                x_min = int(position[2])
                x_max = int(position[3])
                
                # 确保位置在图像范围内
                if img_height > 0 and img_width > 0:
                    y_min = max(0, min(y_min, img_height - 1))
                    y_max = max(0, min(y_max, img_height - 1))
                    x_min = max(0, min(x_min, img_width - 1))
                    x_max = max(0, min(x_max, img_width - 1))
                
                # 检查位置有效性
                if y_min >= y_max or x_min >= x_max:
                    print(f"警告: 符号 {i} 的位置无效 [{y_min},{y_max},{x_min},{x_max}]")
                    continue
                    
                # 保存切片图像
                slice_filename = f"slice_{i}.png"
                slice_path = os.path.join(self.slices_dir, slice_filename)
                
                # 尝试获取图像切片
                if hasattr(self.recognition, 'thImg') and self.recognition.thImg is not None:
                    try:
                        currentFrame = self.recognition.thImg[y_min:y_max+1, x_min:x_max+1]
                        
                        # 检查切片是否为空
                        if currentFrame is None or currentFrame.size == 0:
                            print(f"警告: 符号 {i} 的图像切片为空")
                            continue
                            
                        # 保存切片图像
                        success = cv2.imwrite(slice_path, currentFrame)
                        if not success:
                            print(f"错误: 无法保存符号 {i} 的图像切片到 {slice_path}")
                            continue
                    except Exception as e:
                        print(f"错误: 获取或保存符号 {i} 的图像切片失败: {str(e)}")
                        continue
                else:
                    print(f"警告: 无法获取符号 {i} 的图像切片，二值化图像不存在")
                    continue
                
                # 获取标签
                predicted_label = symbol.get('predicted_label', 'unknown')
                
                # 添加标注信息
                annotation = {
                    "position": [y_min, y_max, x_min, x_max],
                    "predicted_label": predicted_label,
                    "confidence": 0.0,  # 目前Recognition类没有存储置信度
                    "image_slice": slice_path  # 存储切片文件路径
                }
                dataset["annotations"].append(annotation)
                
            except Exception as e:
                print(f"错误: 处理符号 {i} 时发生异常: {str(e)}")
                continue
        
        # 保存数据集文件
        dataset_file = os.path.join(self.output_dir, "dataset.json")
        try:
            with open(dataset_file, 'w') as f:
                json.dump(dataset, f, indent=4, cls=NumpyEncoder)
            print(f"数据集已生成，包含 {len(dataset['annotations'])} 个标注")
        except Exception as e:
            print(f"错误: 保存数据集时出错: {str(e)}")
        
        return self.output_dir